import numpy as np 
from sparc_sim import sparc_sim
from sparc_se import sparc_se
import matplotlib.pyplot as plt
import time

awgn_var      = 1.0                 # AWGN channel noise variance
code_params   = {#'complex':   True,
                 'spatially_coupled': True,
                 #'modulated': True,
                 'P': 16.03,         # Average codeword symbol power constraint
                 'R': 1.6,        # Rate
                 'L': 960,       # Number of sections
                 'M': 128,           # Columns per section
                 #'K': 4,            # MPSK Modulation factor
                 'omega': 6,        # Coupling width
                 'Lambda': 32}       # Coupling length
decode_params = {'t_max': 100}       # Maximum number of iterations
num_of_runs   = 200                  # Number of encoding/decoding trials
rng = np.random.RandomState(seed=None) 

nmse_store  = np.zeros((num_of_runs, decode_params['t_max']))
av_ber = 0.0
for i in range(num_of_runs):
    start_time    = time.perf_counter()
    rng_seed      = rng.randint(2**32-1, size=2, dtype=np.int64).tolist()
    results       = sparc_sim(code_params, decode_params, awgn_var, rng_seed) 
    nmse_store[i] = results['nmse'].mean(axis=1)
    av_ber        = av_ber + results['ber']
    print('Run #{}, BER: {:1.7f}, SER: {:1.4f}, number of iterations: {:3d}, time elapsed: {:2.3f}'
          .format(i, results['ber'],results['ser'], results['t_final'], time.perf_counter()-start_time))
print('Code parameters:', code_params)
av_ber = av_ber / num_of_runs
print(av_ber)